import sys
import os
import torch
from collections import defaultdict
from PIL import Image, ImageDraw, ImageFont

# 载入模型入口，用于生成人脸编码库

# 加载模型所在的模块
# sys.path.append("/Users/mac/PycharmProjects/pytorchTest/face-discern")
# sys.path.append("/Users/mac/PycharmProjects/pytorchTest/face_landmark") # ../face-discern

import fasterrcnn
import face_landmark
import face_discern

# 构建编码数据库使用的文件夹
# 要求图片路径是 {SAMPLES_DIR}/{人物名称}/图片名
SAMPLES_DIR = "face_recognition_samples"

# 输出人脸图片使用的文件夹
FACES_OUTPUT_DIR = "detected_faces"

# 编码数据库的保存路径
# 如果需要重新构建请删除此文件
VECTORS_PATH = "./pt_data/face_recognition_vectors.pt"

# 使用关键点调整脸部范围的次数
ADJUST_FACE_REGION_BY_LANDMARK_TIMES = 3

# 脸部范围相对关键点范围的比例
FACE_REGION_RATIO_BY_LANDMARK = 3.0

# 判断是同一人物需要的分数阈值
FACE_VERIFICATION_THRESHOLD = 0.92

# 用于启用 GPU 支持
DEVICE_TYPE = "cuda" if torch.cuda.is_available() else "cpu"
device = torch.device(DEVICE_TYPE)

def adjust_face_region(image, region, face_landmark_model):
    """使用关键点调整脸部范围"""
    x, y, w, h = region
    child_img = image.crop((x, y, x+w, y+h))
    points = []
    for _ in range(ADJUST_FACE_REGION_BY_LANDMARK_TIMES):
        points = face_landmark_model.detect_landmarks([child_img])[0]
        min_x = min(p[0] for p in points)
        min_y = min(p[1] for p in points)
        max_x = max(p[0] for p in points)
        max_y = max(p[1] for p in points)
        landmark_w = max_x - min_x
        landmark_h = max_y - min_y
        center_x = min_x + landmark_w * 0.5 + x
        center_y = min_y + landmark_h * 0.5 + y
        radius = max(landmark_w, landmark_h) * FACE_REGION_RATIO_BY_LANDMARK / 2
        x0 = int(min(max(center_x - radius, 0), image.size[0]-1))
        x1 = int(min(max(center_x + radius, 0), image.size[0]-1))
        y0 = int(min(max(center_y - radius, 0), image.size[1]-1))
        y1 = int(min(max(center_y + radius, 0), image.size[1]-1))
        x_diff = x0 - x
        y_diff = y0 - y
        x, y, w, h = x0, y0, x1-x0, y1-y0
        points = [ (px - x_diff, py - y_diff) for px, py in points ]
        if w == 0 or h == 0:
            # 检测关键点失败，返回原始范围
            return child_img, region, []
        child_img = image.crop((x, y, x+w, y+h))
    return child_img, (x, y, w, h), points

def build_vectors(faster_rcnn_model, face_landmark_model, face_recognition_model):
    """根据人物和图片构建编码数据库"""
    print("building face recognition vectors from samples")
    names = []
    images = []
    result_names = []
    result_vectors = []
    batch_size = 10
    def process_images():
        # 查找人脸区域
        tensor_in = torch.stack([
            fasterrcnn.image_to_tensor(fasterrcnn.resize_image(img))
            for img in images
        ]).to(device)
        cls_result = faster_rcnn_model(tensor_in)[-1]
        # 截取各图片中的人脸区域
        face_names = []
        face_images = []
        for name, image, result in zip(names, images, cls_result):
            if not result:
                print("no face found for", name)
                continue
            result = fasterrcnn.MyModel.merge_predicted_result(result)
            # 按分数排序并选取分数最高的区域
            result.sort(key=lambda r: -r[2])
            label, box, rpn_score, cls_score = result[0]
            x, y, w, h = fasterrcnn.map_box_to_original_image(box, *image.size)
            if w == 0 or h == 0:
                print("no face found for", name, "due to incorrect size")
                continue
            # 使用脸部关键点调整脸部范围
            child_img, _, points = adjust_face_region(image, (x, y, w, h), face_landmark_model)
            # 保存检测结果用于调试
            if not os.path.isdir(FACES_OUTPUT_DIR):
                os.makedirs(FACES_OUTPUT_DIR)
            face_landmark.draw_points(child_img, points)
            child_img.convert("RGB").save(os.path.join(
                FACES_OUTPUT_DIR, f"{name[0]}_{os.path.splitext(name[1])[0]}.png"))
            face_names.append(name)
            face_images.append(child_img)
        # 转换人脸图片到编码
        tensor_in = torch.stack([
            face_discern.image_to_tensor(face_discern.resize_image(img))
            for img in face_images
        ]).to(device)
        face_vectors = face_recognition_model(tensor_in)
        # 添加名称和编码到结果列表
        for name in face_names:
            result_names.append(name)
        result_vectors.append(face_vectors)
        names.clear()
        images.clear()
    for dirname in os.listdir(SAMPLES_DIR):
        dirpath = os.path.join(SAMPLES_DIR, dirname)
        if not os.path.isdir(dirpath):
            continue
        for filename in os.listdir(dirpath):
            if os.path.splitext(filename)[1].lower() not in [".jpg", ".png", ".jpeg"]:
                continue
            names.append((dirname, filename))
            images.append(Image.open(os.path.join(dirpath, filename)))
            if len(images) >= batch_size:
                process_images()
    if images:
        process_images()
    # 整合编码列表
    result_vector = torch.cat(result_vectors, dim=0)
    # 保存名称和编码
    fasterrcnn.save_tensor((result_names, result_vector), VECTORS_PATH)
    peoples_count = len(set(n[0] for n in result_names))
    print(f"built vectors for {peoples_count} peoples and {result_vector.shape[0]} images")

def recognize_face(
        faster_rcnn_model, face_landmark_model,
        face_recognition_model, face_verification_model,
        image, vectors, top_range=10):
    """识别人脸位置与身份"""
    # 查找人脸区域
    tensor_in = fasterrcnn.image_to_tensor(
        fasterrcnn.resize_image(image)).unsqueeze(0).to(device)
    cls_result = faster_rcnn_model(tensor_in)[-1]
    # 按分数排序并选取分数最高的区域
    result = cls_result[0]
    if not result:
        return None
    result = fasterrcnn.MyModel.merge_predicted_result(result)
    result.sort(key=lambda r: -r[2])
    label, box, rpn_score, cls_score = result[0]
    x, y, w, h = fasterrcnn.map_box_to_original_image(box, *image.size)
    if w == 0 or h == 0:
        return None
    # 使用脸部关键点调整脸部范围
    child_img, (x, y, w, h), points = adjust_face_region(
        image, (x, y, w, h), face_landmark_model)
    # 生成人脸编码
    tensor_in = face_discern.image_to_tensor(
        face_discern.resize_image(child_img)).unsqueeze(0).to(device)
    vector = face_recognition_model(tensor_in)[0]
    # 比较人脸编码，得出最接近的 N 张人脸
    diff = (vector - vectors).pow(2)
    diff_sorted = diff.sum(dim=1).sort()
    indices = diff_sorted.indices[:top_range].tolist()
    # 使用人脸认证模型匹配人脸，身份不明时返回 None
    highest_score = None
    highest_index = None
    scores = face_verification_model(diff[indices]).tolist()
    for index, score in zip(indices, scores):
        if score < FACE_VERIFICATION_THRESHOLD:
            continue
        if highest_score and score <= highest_score:
            continue
        highest_index = index
        highest_score = score
    # 返回人脸图片，人脸范围，人脸关键点，匹配出来的身份，分数
    # 没有匹配到身份时，身份和分数会为 None
    return child_img, (x, y, w, h), points, highest_index, highest_score

def main():
    """组合模型实现人脸检测 + 识别 + 认证"""
    # 让 fasterrcnn 合并区域时不使用 NMS 算法，使得结果区域更大
    fasterrcnn.USE_NMS_ALGORITHM = False

    # 加载 fasterrcnn 模型
    faster_rcnn_model = fasterrcnn.MyModel().to(device) #
    faster_rcnn_model.load_state_dict(fasterrcnn.load_tensor("./pt_data/faster-rcnn-10.5-model.pt"))
    faster_rcnn_model.eval()

    # 加载人脸关键点识别模型
    face_landmark_model = face_landmark.FaceLandmarkModel().to(device)
    face_landmark_model.load_state_dict(fasterrcnn.load_tensor("./pt_data/face-landmark-model.pt"))
    face_landmark_model.eval()

    # 加载人脸识别模型
    face_recognition_model = face_discern.FaceRecognitionModel().to(device)
    face_recognition_model.load_state_dict(face_discern.load_tensor(
        "./pt_data/face-recognition-model.recognition.pt"))
    face_recognition_model.eval()

    # 加载人脸验证模型
    face_verification_model = face_discern.FaceVerificationModel().to(device)
    face_verification_model.load_state_dict(face_discern.load_tensor(
        "./pt_data/face-recognition-model.verification.pt"))
    face_verification_model.eval()

    # 根据人物和图片构建编码数据库
    if not os.path.isfile(VECTORS_PATH):
        build_vectors(faster_rcnn_model, face_landmark_model, face_recognition_model)

    # 读取编码数据库
    names, vectors = fasterrcnn.load_tensor(VECTORS_PATH)

    # 询问图片路径，并识别人脸
    ### 修改
    font = ImageFont.truetype("/Library/Fonts/Arial Unicode.ttf", 20) # wqy-microhei.ttc

    while True:
        try:
            # 打开图片
            image_path = input("Image path: ")
            if not image_path:
                continue
            image = Image.open(image_path)
            # 识别人脸
            result = recognize_face(
                faster_rcnn_model, face_landmark_model,
                face_recognition_model, face_verification_model, image, vectors)
            if not result:
                print("no face found\n")
                continue
            # 根据结果标记到图片上
            child_img, (x, y, w, h), points, index, score = result
            points = [ (px + x, py + y) for px, py in points ]
            label = f"{names[index][0]}: {score:.2f}" if index else "Unknown"
            print("result:", label)
            color = "#FF0000"
            draw = ImageDraw.Draw(image)
            draw.rectangle((x, y, x+w, y+h), outline=color, width=3)
            draw.text((x, y-20), label, fill=color, font=font)
            face_landmark.draw_points(image, points)
            image.save("img_output.png")
            print("saved to img_output.png")
            print()
        except Exception as e:
            print("error:", e)

if __name__ == "__main__":
    main()
